import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from configs import DEVICE, EPOCHS, INPUT_SIZE, HIDDEN_SIZE1, HIDDEN_SIZE2, OUTPUT_SIZE, MODEL_NAME
from model import FCNet, CNNet
from dataloader import prepare_test_data

def test_model(model,device,test_loader):
    model.eval()
    model.load_state_dict(torch.load('./ckpt/'+MODEL_NAME+str(EPOCHS)+'.pth'))
    correct = 0
    test_loss = 0
    with torch.no_grad():
        for data,label in test_loader:
            if MODEL_NAME == 'FCN' :
                data, label = data.reshape(-1,28*28).to(device), label.to(device)
            elif MODEL_NAME == 'CNN':
                data, label = data.to(device), label.to(device)
            output = model(data)
            test_loss += F.cross_entropy(output,label).item()
            predict = torch.max(torch.tensor(output),dim=1).indices
            correct += predict.eq(label.view_as(predict)).sum().item()
        test_loss /= len(test_loader.dataset)
        correct /= len(test_loader.dataset)
        print("Test--Average loss:{:.4f},Accuracy:{:.3f}\n".format(test_loss,100.0*correct))
if MODEL_NAME == 'FCN' :
    model = FCNet(INPUT_SIZE, HIDDEN_SIZE1, HIDDEN_SIZE2, OUTPUT_SIZE).to(DEVICE)
elif MODEL_NAME == 'CNN' :
    model = CNNet(INPUT_SIZE, OUTPUT_SIZE).to(DEVICE)
test_loader=prepare_test_data()
test_model(model,DEVICE,test_loader)